import { useContext, useEffect, useRef } from "react";
import { IsTrainingContext } from "./context";
import * as echarts from "echarts";
import { WEBSOCKET_URL } from "../../api";
import terminalIcon from "../images/terminal_light.svg";
import chartIcon from "../images/Chart_alt_light.svg";
import styles from "./styles.module.scss";
import { ChartData, getOption } from "../../util";

const DATA_FOR_SHOW = 1000;
function TrainMonitor() {
  return (
    <div className={styles.monitor}>
      <TrainLog></TrainLog>
      <TrainChart></TrainChart>
    </div>
  );
}

function TrainLog() {
  const { isTraining, setIsTraining } = useContext(IsTrainingContext);
  const ref = useRef<HTMLDivElement>(document.createElement("div"));

  // 初始化数据
  useEffect(() => {
    if (isTraining) {
      ref.current.innerHTML = "";
    }
  }, [isTraining]);

  useEffect(() => {
    const ws = new WebSocket(WEBSOCKET_URL + "?task=train&monitor=log");
    ws.onmessage = (e) => {
      const msg = e.data;
      if (msg.startsWith("exit")) {
        setIsTraining(false);
        ref.current.appendChild(document.createTextNode("训练进程退出"));
        ref.current.scrollTop = ref.current.scrollHeight;
        return;
      }
      const text = document.createTextNode(msg);
      const br = document.createElement("br");
      ref.current.appendChild(text);
      ref.current.appendChild(br);
      if (ref.current.childNodes.length >= DATA_FOR_SHOW * 2) {
        ref.current.removeChild(ref.current.firstChild!);
        ref.current.removeChild(ref.current.firstChild!);
      }
      ref.current.scrollTop = ref.current.scrollHeight;
    };
    return () => {
      ws.close();
    };
  }, [setIsTraining]);

  return (
    <div className={`${styles.logArea} ${styles.myCard}`}>
      <h5>
        <img alt="" src={terminalIcon} />
        训练日志
      </h5>
      <div className={styles.msgs} ref={ref}></div>
    </div>
  );
}

function TrainChart() {
  const { isTraining, setIsTraining } = useContext(IsTrainingContext);
  const ref = useRef<HTMLDivElement>(document.createElement("div"));
  let chartData = useRef<ChartData>({ x: [], y: [], avg: [] });

  //初始化数据
  useEffect(() => {
    if (isTraining) {
      chartData.current = { x: [], y: [], avg: [] };
    }
  }, [isTraining]);

  // 获取chart数据
  useEffect(() => {
    const ws = new WebSocket(WEBSOCKET_URL + "?task=train&monitor=chart");
    const chart = echarts.init(ref.current);
    ws.onmessage = (e) => {
      const info = e.data;
      const json = JSON.parse(info);
      chartData.current.x.push(json.x);
      chartData.current.y.push(json.y);
      chartData.current.avg.push(json.avg);
      if (chartData.current.x.length > DATA_FOR_SHOW) {
        chartData.current.x.shift();
        chartData.current.y.shift();
        chartData.current.avg.shift();
      }
      chart.setOption(getOption(chartData.current));
    };
    return () => {
      chart.dispose();
      ws.close();
    };
  }, [setIsTraining]);

  return (
    <div className={`${styles.chartArea} ${styles.myCard}`}>
      <h5>
        <img alt="" src={chartIcon} />
        训练进度
      </h5>
      <div className={styles.chart} ref={ref}></div>
    </div>
  );
}

export default TrainMonitor;
